import sys
import numpy as np
from PyQt6.QtWidgets import (
    QApplication, QMainWindow, QVBoxLayout, QPushButton, QTextEdit, QLabel, QWidget, QTableWidget, QTableWidgetItem
)
from PyQt6.QtGui import QPixmap
import matplotlib.pyplot as plt
import networkx as nx
from matplotlib.patches import Rectangle
from matplotlib.transforms import Affine2D
from matplotlib.patches import FancyArrowPatch, Circle
from Core import *


class GraphVisualizer(QWidget):
    def __init__(self):
        super().__init__()
        self.setWindowTitle("Graph Network Visualizer")
        self.setGeometry(100, 100, 900, 700)

        # Main layout
        self.layout = QVBoxLayout(self)

        # Add widgets
        self.calculate_button = QPushButton("Generate Graph and Matrix")
        self.calculate_button.clicked.connect(self.calculate_graph)
        self.layout.addWidget(self.calculate_button)

        self.result_text = QTextEdit()
        self.result_text.setReadOnly(True)
        self.layout.addWidget(self.result_text)

        self.graph_label = QLabel("Graph Visualization:")
        self.layout.addWidget(self.graph_label)
        self.graph_pixmap = QLabel()
        self.layout.addWidget(self.graph_pixmap)

        self.table_label = QLabel("Adjacency Matrix:")
        self.layout.addWidget(self.table_label)
        self.adjacency_table = QTableWidget()
        self.layout.addWidget(self.adjacency_table)
        self.collision_vector = None
        self.vector_list = None
        self.adjacency_list = None
        self.adjacency_matrix = None
        self.min_avg_weight = None
        self.best_cycle = None

    def calculate_graph(self):
        try:
            # Placeholder collision vector
            self.collision_vector = np.array([1, 0, 1, 1, 0, 1, 0])

            # Generate graph
            self.vector_list, self.adjacency_list = generate_adjacency_list(self.collision_vector, max_iterations=10)
            self.adjacency_matrix = adjacency_list_to_matrix(self.adjacency_list, len(self.vector_list))
            self.min_avg_weight, self.best_cycle = find_minimum_average_cycle(self.adjacency_matrix)

            # Update UI with results
            self.result_text.append(f"Generated {len(vector_list)} vectors.")
            self.result_text.append(f"Minimum Average Weight Cycle: {min_avg_weight}")
            self.result_text.append(f"Best Cycle: {best_cycle}")

            self.update_adjacency_table(adjacency_matrix)
            self.visualize_graph(adjacency_list,vector_list)
        except Exception as e:
            self.result_text.append(f"Error: {e}")

    def update_adjacency_table(self, adjacency_matrix):
        num_nodes = len(adjacency_matrix)
        self.adjacency_table.setRowCount(num_nodes)
        self.adjacency_table.setColumnCount(num_nodes)
        for i in range(num_nodes):
            for j in range(num_nodes):
                item = QTableWidgetItem(f"{adjacency_matrix[i][j]:.1f}" if adjacency_matrix[i][j] != np.inf else "∞")
                self.adjacency_table.setItem(i, j, item)


    def visualize_graph(self, adjacency_list, vector_list):
        # 使用 MultiDiGraph 以支持多重边
        graph = nx.MultiDiGraph()

        # 构建图
        for src, edges in adjacency_list.items():
            for dst, weight in edges:
                graph.add_edge(src, dst, weight=weight)

        # 定义布局
        pos = nx.spring_layout(graph, seed=42)  # 设置种子确保布局一致

        plt.figure(figsize=(12,7))
        ax = plt.gca()


        # 初始化偏移字典，记录每对节点的标签偏移
        label_offsets = {}

        # 绘制边和标签
        for u, v, k, d in graph.edges(keys=True, data=True):
            is_reverse = False
            is_loop = False
            # 判断是否为反向边
            if (v, u) in label_offsets:
                is_reverse = True
            if u == v:
                is_reverse = True
                is_loop = True

            # 动态调整弧度，反向边弧度取负
            rad = 0.2 + k * 0.2

            # 绘制边
            if not is_loop:
                nx.draw_networkx_edges(
                    graph,
                    pos,
                    edgelist=[(u, v)],
                    connectionstyle=f"arc3,rad={rad}",
                    arrowstyle="->",
                    arrowsize=15,
                )
                nx.draw_networkx_edges(
                    graph,
                    pos,
                    edgelist=[(u, v)],
                    connectionstyle=f"arc3,rad={rad}",
                    arrowstyle="->",
                    arrowsize=15,
                    min_source_margin=100,  # 增加源节点边缘的距离
                    min_target_margin=100  # 增加目标节点边缘的距离
                )
            else:
                # 自环：绘制圆
                center_x, center_y = pos[u]
                center_y+=0.005
                radius = rad/5  # 圆的半径
                loop_circle = Circle(
                    (center_x, center_y+radius),
                    radius,
                    color="black",
                    fill=False,
                    # lw=1.5,
                )
                ax.add_patch(loop_circle)

                # 绘制自环箭头
                arrow_x = center_x + radius
                arrow_y = center_y + radius
                arrow = FancyArrowPatch(
                    (arrow_x , arrow_y - 0.001),
                    (arrow_x , arrow_y + 0.001),
                    arrowstyle="->",
                    color="black",
                    mutation_scale=15,
                )
                ax.add_patch(arrow)

            # 初始化或更新标签偏移值
            if (u, v) not in label_offsets:
                label_offsets[(u, v)] = []
            label_offsets[(u, v)].append(k)

            # 获取标签偏移位置
            x_middle = (pos[u][0] + pos[v][0]) / 2
            y_middle = (pos[u][1] + pos[v][1]) / 2

            # 偏移随 k 递增，反向边标签偏移方向相反
            offset = 0.05 * (k*0.7 + 1)
            if not is_reverse:
                offset *= -1

            # 绘制权值标签
            plt.text(
                x_middle + offset,
                y_middle + offset,
                f"{d['weight']}",
                fontsize=12,
                color="red",
                ha="center",
                va="center",
            )
            print(u,v,k,d,is_reverse,is_loop)
        # 绘制节点矩形
        for node, (x, y) in pos.items():
            vector = vector_list[node]
            text = str(vector)
            width = len(text) * 0.02  # 根据内容长度动态调整宽度
            height = 0.02  # 固定高度

            # 绘制矩形
            rect = Rectangle(
                (x - width / 2, y - height / 2),
                width,
                height,
                color="skyblue",
                ec="black",
                lw=1.5,

            )
            ax.add_patch(rect)

            # 绘制节点标签
            plt.text(
                x,
                y,
                text,
                fontsize=8,
                ha="center",
                va="center",
            )

        # 保存并更新 QLabel
        plt.savefig("graph.png", format="png")
        plt.close()

        pixmap = QPixmap("../graph.png")
        self.graph_pixmap.setPixmap(pixmap)


if __name__ == "__main__":
    app = QApplication(sys.argv)
    window = GraphVisualizer()
    window.show()
    sys.exit(app.exec())
